# 1、支持无人机飞直线和斜线的真实时间对齐
# 2、支持位置冲突/边冲突 的安全时段SLACK检测，并提供时段合并以降低冲突检测时候遍历的已占据时段数量
import numpy as np
import math
import heapq

import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from collections import defaultdict
from datetime import datetime

from algorithm.Astar import AStarPlanner

# ====== 配置区域 ======
TIME_SCALE = 100      # 时间离散比例，将真实时间*TIME_SCALE转为整数刻度
SLACK = 2           # 每格时间缓冲区间，单位秒
SUBSTEPS = 5

# ====== 新增：辅助函数，用于合并重叠区间 ======
def merge_intervals(intervals):  # 添加
    if not intervals:
        return []
    intervals.sort(key=lambda x: x[0])  # 修改：按开始时间排序
    merged = [list(intervals[0])]
    for start, end in intervals[1:]:
        if start <= merged[-1][1]:  # 修改：检测重叠或相接
            merged[-1][1] = max(merged[-1][1], end)  # 修改：合并区间
        else:
            merged.append([start, end])  # 添加：新建不重叠区间
    return [(s, e) for s, e in merged]  # 添加

class ReservationTable:
    """时空预留表，记录格点和边的占用时间区间"""
    def __init__(self):
        # 存储格点和边的占用区间
        self.vertex_res = defaultdict(list)        # key: (x,y) -> [(start,end), ...]
        self.edge_res   = defaultdict(list)        # key: ((x0,y0),(x1,y1)) -> [(start,end), ...]

    def reserve_path(self, path):
        # path: [(x, y, raw_t), ...]
        for (x, y, raw_t) in path:
            t = raw_t * TIME_SCALE
            interval = t - SLACK*TIME_SCALE, t + SLACK*TIME_SCALE
            self.vertex_res[(x, y)].append(interval)
        for ((x0, y0, t0), (x1, y1, t1)) in zip(path, path[1:]):
            start = t0 * TIME_SCALE - SLACK*TIME_SCALE
            end   = t1 * TIME_SCALE + SLACK*TIME_SCALE
            interval = (start, end)
            # 正向与反向边都标记占用
            key_fwd = ((x0, y0), (x1, y1))
            key_rev = ((x1, y1), (x0, y0))
            self.edge_res[key_fwd].append(interval)
            self.edge_res[key_rev].append(interval)

        # ======= 合并顶点和边的所有区间，减少冗余 ======
        for key in list(self.vertex_res.keys()):  # 添加
            self.vertex_res[key] = merge_intervals(self.vertex_res[key])  
        for key in list(self.edge_res.keys()): 
            self.edge_res[key] = merge_intervals(self.edge_res[key]) 


    def vertex_conflict(self, x, y, t):
        for (s, e) in self.vertex_res[(x, y)]:
            # 时间区间重叠即冲突
            if t > s and t < e:
                return True
        return False

    def edge_conflict(self, x0, y0, x1, y1, t_enter, t_exit):
        # 检查正、反向边冲突
        for key in [((x0, y0), (x1, y1)), ((x1, y1), (x0, y0))]:
            for (s, e) in self.edge_res.get(key, []):
                # if t_enter < e and t_exit > s:
                if not (t_exit < s or t_enter > e):
                    return True
        # 检查对角穿越的双向冲突
        dx, dy = x1 - x0, y1 - y0
        if abs(dx) == 1 and abs(dy) == 1:
            # 交叉边两个方向
            cross_edges = [((x0, y1), (x1, y0)), ((x1, y0), (x0, y1))]
            for key in cross_edges:
                for (s, e) in self.edge_res.get(key, []):
                    # if t_enter < e and t_exit > s:
                    if not (t_exit < s or t_enter > e):
                        return True
        return False

class STNode:
    # __slots__ 告诉 Python 这个类 只 会有这六个属性：
    # f = g + h
    __slots__ = ('x','y','t','g','f','parent')
    def __init__(self, x, y, t, g, f, parent):
        self.x, self.y, self.t = x, y, t
        self.g, self.f = g, f
        self.parent = parent
        # 重载（override）小于比较运算 (<)，让两个 STNode 能直接用 node1 < node2 来比较它们的 f 值。确保在 heapq 中能够按 f 值正确排序
    def __lt__(self, other):
        return self.f < other.f

class SpaceTimeAStar:
    """在 (x, y, t) 空间上做 A*，避开静/动态障碍 + 边冲突"""
    def __init__(self, planner:AStarPlanner, 
                 reservation:ReservationTable, 
                 speed,
                 max_t=500):
        self.planner = planner
        self.resoluiton = planner.resolution
        self.RT      = reservation
        # 无人机速度
        self.speed = speed
        # 地图信息来自原 AStarPlanner
        self.obs_map = planner.obstacle_map
        self.xw, self.yw = planner.x_width, planner.y_width
        self.min_x, self.min_y = planner.min_x, planner.min_y
        # 取原 motion model，加上 wait 动作
        # self.motion = planner.motion + [[0, 0, 1.0]]
        # 根据栅格分辨率和飞行速度动态生成 motion
        d_straight = self.resoluiton
        d_diag     = round(self.resoluiton * math.sqrt(2), 3)
        t_straight = d_straight / self.speed
        t_diag     = d_diag / self.speed
        self.motion = [
            ( 1,  0, t_straight),
            (-1,  0, t_straight),
            ( 0,  1, t_straight),
            ( 0, -1, t_straight),
            ( 1,  1, t_diag),
            ( 1, -1, t_diag),
            (-1,  1, t_diag),
            (-1, -1, t_diag),
            (0, 0, t_straight),
        ]
        self.max_t  = max_t


    def heuristic(self, x, y, gx, gy):
        # 欧氏距离启发
        dist = math.hypot(gx - x, gy - y) * self.resoluiton
        t_min = dist/self.speed

        return t_min

    def plan(self, sx, sy, gx, gy):
        # 把世界坐标转换到网格索引
        sx_i = self.planner.calc_xy_index(sx, self.min_x)
        sy_i = self.planner.calc_xy_index(sy, self.min_y)
        gx_i = self.planner.calc_xy_index(gx, self.min_x)
        gy_i = self.planner.calc_xy_index(gy, self.min_y)
        # closed
        closed_set = dict()  # key:(x,y,t) -> value: g
        open_set = dict()
        h0 = self.heuristic(sx_i, sy_i, gx_i, gy_i)
        # 起点的父节点为空
        start = STNode(sx_i, sy_i, 0, 0.0, h0, None)

        """
        方法一：和Astar一样的节点弹出与验证方式
        
        """
        key_start = (start.x, start.y, start.t)
        open_set[key_start] = start
        # 主循环
        while True:
            if len(open_set) == 0:
                print("Open set is empty..")
                break
            #print('plan exe')
            cur_key = min(
                open_set,
                key=lambda o: open_set[o].g + self.heuristic(open_set[o].x, open_set[o].y, gx_i, gy_i))
            current = open_set[cur_key]

            # 到达目标
            if current.x == gx_i and current.y == gy_i:
                # 回溯
                path = []
                node = current
                while node:
                    # 返回的path是栅格索引而不是真实坐标
                    path.append((node.x, node.y, node.t))
                    node = node.parent
                return list(reversed(path))
            
            # Remove the item from the open set
            del open_set[cur_key]

            # Add it to the closed set
            closed_set[cur_key] = current

            # 扩展动作
            for dx, dy, dt_cost in self.motion:
                nx, ny = current.x + dx, current.y + dy
                raw_nt     = current.t + dt_cost  # 步长和扩展栅格有关（直/斜
                # 处理浮点数，有利于后续加入closed后仍被识别（浮点数每次计算不一致）
                raw_nt = round(raw_nt, 2)
                key2  = (nx, ny, raw_nt)

                # 对等待步进行必要性检测
                if dx==0 and dy==0:
                    # 先检测非等待的8个方向，看有没有至少一个可走
                    blocked = True
                    for ddx, ddy, dd_cost in self.motion:
                        if (ddx, ddy) == (0,0): continue
                        tx, ty = current.x + ddx, current.y + ddy
                        ttest  = current.t + dd_cost
                        itt    = int(ttest * TIME_SCALE)
                        if (0 <= tx < self.xw and 0 <= ty < self.yw
                            and not self.obs_map[tx][ty]
                            and key2 not in closed_set
                            and not self.RT.vertex_conflict(tx, ty, itt)
                            and not self.RT.edge_conflict(current.x, current.y, tx, ty, 
                                                        int(current.t*TIME_SCALE), itt)):
                            blocked = False
                            break
                    # 如果至少有一个方向可走，就跳过等待
                    if not blocked:
                        continue


                # 乘上时间尺度方便进行比较100/144
                nt = int(raw_nt * TIME_SCALE)
                # 用时间定义代价g
                g2 = current.g + dt_cost
                # 用欧几里得距离作为引导
                h2 = self.heuristic(nx, ny, gx_i, gy_i)
                # 每次扩展新节点时候会算一下f = g + h
                node2 = STNode(nx, ny, raw_nt, g2, g2 + h2, current)


                if nt > self.max_t*TIME_SCALE:
                    continue
                # 约束1：边界 & 静态障碍
                if not (0 <= nx < self.xw and 0 <= ny < self.yw): 
                    continue
                # 约束2：无障碍物
                if self.obs_map[nx][ny]:
                    continue
                # 时空冲突：顶点 & 边
                if self.RT.vertex_conflict(nx, ny, nt):
                    continue
                # 当前点时间
                prev_t_scaled = int(current.t * TIME_SCALE)
                if self.RT.edge_conflict(current.x, current.y, nx, ny, prev_t_scaled, nt):
                    continue
                # 约束3：不扩展已经扩展的
                if key2 in closed_set:
                    continue
                
                if key2 not in open_set:
                    open_set[key2] = node2  # discovered a new node
                else:
                    if open_set[key2].g > node2.g:
                        # This path is the best until now. record it
                        open_set[key2] = node2

        # 如果找不到
        raise RuntimeError("Space–Time A* failed to find a path")


        """
        方法二hash 表弹出节点，可能出现不能将更小代价节点更换的情况
        """
        # open_heap：优先队列，用 (f, node) 保证每次弹出 f 最小的节点。
        # closed：哈希表，避免对同一(x,y,t) 重复处理代价更高的节点。
        # open_heap = []
        # # 建立小跟堆
        # heapq.heappush(open_heap, (start.f, start))
        # heapq.heappush(open_heap, item) 会将 item 插入到 open_heap，并维护它的堆性质；
        """
        这里的 item 是一个二元组 (start.f, start):
        第一个元素 start.f(即 g+h)用来比较大小,堆顶始终是 f 最小的；
        第二个元素 start 是你的节点对象，当 f 相同时,Python 会去比较 start 对象本身（此时又会调用你在 STNode 里定义的 __lt__ 方法）；
        """
        # # 主循环
        # while open_heap:

        #     # 弹出open_heap中f最小的元素
        #     _, cur = heapq.heappop(open_heap)
        #     key = (cur.x, cur.y, cur.t)
        #     # 把当前节点的时空三元组 (x,y,t) 做成 key，若已在 closed（闭集）中，说明之前以更优代价到过这里，就跳过
        #     if key in closed_set:
        #         continue
        #     closed_set[key] = cur.g

        #     # 到达目标
        #     if cur.x == gx_i and cur.y == gy_i:
        #         # 回溯
        #         path = []
        #         node = cur
        #         while node:
        #             # 返回的path是栅格索引而不是真实坐标
        #             path.append((node.x, node.y, node.t))
        #             node = node.parent
        #         return list(reversed(path))

        #     # 扩展动作
        #     for dx, dy, dt_cost in self.motion:
        #         nx, ny = cur.x + dx, cur.y + dy
        #         nt     = cur.t + 1  # 固定步长为 1
        #         if nt > self.max_t:
        #             continue
        #         # 边界 & 静态障碍
        #         if not (0 <= nx < self.xw and 0 <= ny < self.yw): 
        #             continue
        #         if self.obs_map[nx][ny]:
        #             continue
        #         # 时空冲突：顶点 & 边
        #         if self.RT.vertex_conflict(nx, ny, nt):
        #             continue
        #         if self.RT.edge_conflict(cur.x, cur.y, nx, ny, cur.t):
        #             continue
        #         # 用时间定义代价g
        #         g2 = cur.g + dt_cost
        #         # 用欧几里得距离作为引导
        #         h2 = self.heuristic(nx, ny, gx_i, gy_i)
        #         # 每次扩展新节点时候会算一下f = g + h
        #         node2 = STNode(nx, ny, nt, g2, g2 + h2, cur)
        #         key2  = (nx, ny, nt)
        #         if key2 in closed_set and closed_set[key2] <= g2:
        #             continue
        #         # 新节点添加进堆
        #         heapq.heappush(open_heap, (node2.f, node2))

        # # 如果找不到2.0
        # raise RuntimeError("Space–Time A* failed to find a path")

        

def multi_uav_planning(uav_list, planner, speed, max_t=500, existing_reservation_table=None):
    """
    uav_list: [(sx,sy,gx,gy), ...] 世界坐标[m]
    返回: 每架 UAV 的时空路径 list of [(x,y,t), ...]
    """
    if existing_reservation_table is not None:
        reservation = existing_reservation_table
    else:
        reservation = ReservationTable()
    st_planner  = SpaceTimeAStar(planner, reservation, speed, max_t)
    all_paths   = []
    # 用时空规划器进行规划
    print(f"Planning space-time A*...")
    for i, (sx, sy, gx, gy) in enumerate(uav_list):
        path = st_planner.plan(sx, sy, gx, gy)
        # 将已有的路径存储到时空预留表
        reservation.reserve_path(path)
        all_paths.append(path)
    print(f"Get space-time A* Path")
    return all_paths


def plot_static_trajectories(all_time_paths, planner, uav_ids=None, entry_points=None):
    """
    Draws static trajectories for all UAVs on a single plot.
    all_time_paths: list of paths, each path is [(x_idx, y_idx, t), ...]
    planner: instance of AStarPlanner (for coordinate conversion)
    uav_ids: list of actual UAV IDs to display in legend
    entry_points: dict of {uav_id: (x_idx, y_idx)} for patrol entry points
    """
    n = len(all_time_paths)
    colors = plt.get_cmap('tab10')(range(n))
    fig, ax = plt.subplots(figsize=(10, 10))

    # 绘制障碍物
    ox, oy = [], []
    for i in range(planner.x_width):
        for j in range(planner.y_width):
            if planner.obstacle_map[i][j]:
                ox.append(planner.calc_grid_position(i, planner.min_x))
                oy.append(planner.calc_grid_position(j, planner.min_y))
    ax.scatter(ox, oy, c='k', s=5, label='Obstacles')
    ax.set_aspect('equal')

    # 绘制轨迹和关键点
    for idx, path in enumerate(all_time_paths):
        uav_id = uav_ids[idx] if uav_ids and idx < len(uav_ids) else (idx + 1)
        xs = [planner.calc_grid_position(xi, planner.min_x) for xi, _, _ in path]
        ys = [planner.calc_grid_position(yi, planner.min_y) for _, yi, _ in path]
        
        # 绘制轨迹线
        label = f'UAV {uav_id}'
        ax.plot(xs, ys, '-', linewidth=2, color=colors[idx], label=label)
        
        # 绘制起点（圆圈）
        ax.scatter(xs[0], ys[0], marker='o', s=50, color=colors[idx], 
                  edgecolors='k', linewidths=2, label=f'UAV {uav_id} Start' if idx == 0 else "")
        
        # 绘制终点（X标记）
        ax.scatter(xs[-1], ys[-1], marker='X', s=50, color=colors[idx], 
                  edgecolors='k', linewidths=2, label=f'UAV {uav_id} End' if idx == 0 else "")
        
        # 绘制入口点（菱形标记）- 仅对巡航任务无人机
        if entry_points and uav_id in entry_points:
            entry_x_idx, entry_y_idx = entry_points[uav_id]
            entry_x = planner.calc_grid_position(entry_x_idx, planner.min_x)
            entry_y = planner.calc_grid_position(entry_y_idx, planner.min_y)
            ax.scatter(entry_x, entry_y, marker='D', s=0, color=colors[idx], 
                      edgecolors='k', linewidths=2, label=f'UAV {uav_id} Entry' if idx == 0 else "")

    # 添加图例和标签
    ax.set_title('UAV Trajectories with Patrol Entry Points')
    ax.set_xlabel('X [m]')
    ax.set_ylabel('Y [m]')
    
    # 整理图例，避免重复
    handles, labels = ax.get_legend_handles_labels()
    by_label = dict(zip(labels, handles))
    ax.legend(by_label.values(), by_label.keys(), title='Legend')
    
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.show()



def animate_uav_trajectories(all_time_paths, planner, interval=100):
    """
    Creates a smooth animation of UAVs moving over time with interpolation.
    Returns a FuncAnimation object.
    """
    # build keyframes per UAV
    keyframes = []  # list of dicts: {frame: (x, y)}
    frame_times = set()
    for path in all_time_paths:
        d = {}
        for xi, yi, t in path:
            f = t * TIME_SCALE
            xw = planner.calc_grid_position(xi, planner.min_x)
            yw = planner.calc_grid_position(yi, planner.min_y)
            d[f] = (xw, yw)
            frame_times.add(f)
        keyframes.append(d)

    # sorted unique keyframe times
    base_frames = sorted(frame_times)
    # create smooth frame timeline with interpolation
    smooth_times = []
    for i in range(len(base_frames)-1):
        t0, t1 = base_frames[i], base_frames[i+1]
        dt = (t1 - t0) / SUBSTEPS
        for k in range(SUBSTEPS): smooth_times.append(t0 + k*dt)
    smooth_times.append(base_frames[-1])

    # precompute interpolated positions for each UAV
    n = len(keyframes)
    uav_smooth = []
    for d in keyframes:
        times = np.array(sorted(d.keys()))
        xs = np.array([d[t][0] for t in times])
        ys = np.array([d[t][1] for t in times])
        # linear interpolation functions
        xi = np.interp(smooth_times, times, xs)
        yi = np.interp(smooth_times, times, ys)
        uav_smooth.append((xi, yi))

    # setup plot
    fig, ax = plt.subplots(figsize=(8,8))
    ox, oy = [], []
    for i in range(planner.x_width):
        for j in range(planner.y_width):
            if planner.obstacle_map[i][j]:
                ox.append(planner.calc_grid_position(i, planner.min_x))
                oy.append(planner.calc_grid_position(j, planner.min_y))
    ax.scatter(ox, oy, c='k', s=5)
    ax.set_aspect('equal')

    # limits
    all_x = np.concatenate([pos[0] for pos in uav_smooth])
    all_y = np.concatenate([pos[1] for pos in uav_smooth])
    ax.set_xlim(all_x.min()-1, all_x.max()+1)
    ax.set_ylim(all_y.min()-1, all_y.max()+1)
    ax.set_title('Smooth UAV Trajectory Animation')

    # init artists
    colors = plt.get_cmap('tab10')(range(n))
    scatters = [ax.scatter([], [], s=80, color=colors[i]) for i in range(n)]
    trails = [ax.plot([], [], linewidth=1, color=colors[i], alpha=0.7)[0] for i in range(n)]

    def update(i):
        for idx, (xi, yi) in enumerate(uav_smooth):
            # current position
            x, y = xi[i], yi[i]
            scatters[idx].set_offsets([[x, y]])
            # trail
            trails[idx].set_data(xi[:i+1], yi[:i+1])
        ax.set_xlabel(f'Time: {smooth_times[i]/TIME_SCALE:.2f}s')
        return scatters + trails

    anim = FuncAnimation(fig, update, frames=len(smooth_times), interval=interval, blit=False)
    return anim



if __name__ == "__main__":
    # —— 与你 main() 中示例一致的参数 —— 
    grid_size, robot_radius = 5.0, 1.0

    ox, oy = [], []
    # for i in range(-10, 60):
    #     ox.append(i);    oy.append(-10.0)
    # for i in range(-10, 60):
    #     ox.append(60.0); oy.append(i)
    # for i in range(-10, 61):
    #     ox.append(i);    oy.append(60.0)
    # for i in range(-10, 61):
    #     ox.append(-10.0);oy.append(i)
    # for i in range(-10, 40):
    #     ox.append(20.0); oy.append(i)
    # for i in range(0, 40):
    #     ox.append(40.0); oy.append(60.0 - i)
    # 边界障碍物 ,small square
    ox.extend(range(-100, 101))
    oy.extend([-100] * 201)#x 从 -100 到 100，y 固定为 -100
    ox.extend(range(-100, 101))
    oy.extend([100] * 201)
    oy.extend(range(-100, 101))
    ox.extend([-100] * 201)
    oy.extend(range(-100, 101))
    ox.extend([100] * 201)
    
    # 初始化单机 A*
    planner = AStarPlanner(ox, oy, grid_size, robot_radius)

    # 定义多 UAV 起终点
    # 起降点8相对于1的坐标: 北向33.26952852米, 东向12.10452645米
    # 起降点2相对于1的坐标: 北向5.15945072米, 东向-11.77504894米
    # 起降点3相对于1的坐标: 北向9.41821120米, 东向4.54526410米
    # 起降点7相对于1的坐标: 北向28.08784494米, 东向12.10453198米

    uav_list = [
        # (0.0, 0.0, 289.50024993, 310.57244424),#1>>any
        # (9.41821120, 4.54526410, 90.34827360, 232.95038901),#3>>any
        # (0.0, 0.0, 289.50024993, 310.57244424),#1>>any
        # (9.41821120, 4.54526410, 90.34827360, 232.95038901),#3>>any

        # (5.15945072, -11.77504894, 28.08784494, 12.10453198),#2>>7
        # (-2.87, 1.70, 28.08784494, 12.10453198), #any>>7
        # (40.0, 10.0, 10.0, 30.0),

        # (8.0, 10.0, 50.0, 40.0),
        # (6.0, 10.0, 50.0, 30.0),
        # (2.0, 10.0, 50.0, 20.0),
        # (0.0, 10.0, 50.0, 10.0),
        # (0.0, 20.0, 55.0, 20.0),
        # (0.0, 30.0, 55.0, 30.0),
        # … 如果需要更多 UAV，继续追加
    ]

    start_time = datetime.now() 
    # 进行时空 A* 规划
    all_time_paths = multi_uav_planning(uav_list, planner, speed=3, max_t=300)

    end_time = datetime.now()
    elapsed_time = end_time - start_time
    print(f"代码执行时间：{elapsed_time.total_seconds():.4f} 秒")  # 输出：代码执行时间：0.0321 秒
    # 打印每条路径（示例）
    for uav_idx, path in enumerate(all_time_paths, 1):
        print(f"\nUAV {uav_idx} 路径（共 {len(path)} 个点）:")
        for point in path:
            x_idx, y_idx, t = point
            # 转换为真实世界坐标（米）
            x_real = planner.calc_grid_position(x_idx, planner.min_x)
            y_real = planner.calc_grid_position(y_idx, planner.min_y)
            print(f"  坐标: ({x_real:.2f}, {y_real:.2f}), 时间: {t:.2f} s")    
    # for i, path in enumerate(all_time_paths):
    #     print(f"UAV{i+1} path (x,y,t): {path}")

    plot_static_trajectories(all_time_paths, planner,uav_ids=[1,2])  # 传入实际 UAV ID 列表

    anim = animate_uav_trajectories(all_time_paths, planner, interval=500)
    plt.show()  # 或者 anim.save('uav_anim.mp4')

